import torch
import torch.nn as nn
from mmcls.models import CLASSIFIERS, ImageClassifier
from mmcls.utils import get_root_logger


@CLASSIFIERS.register_module()
class DLAClassifier(ImageClassifier):
    def __init__(self, backbone, neck=None, head=None, pretrained=None):
        super(DLAClassifier, self).__init__(backbone, neck, head, pretrained)

    def init_weights(self, pretrained):
        super(ImageClassifier, self).init_weights(pretrained)
        self.backbone.init_weights(pretrained=pretrained)
        if self.with_neck:
            if isinstance(self.neck, nn.Sequential):
                for m in self.neck:
                    m.init_weights()
            else:
                self.neck.init_weights()
        if self.with_head:
            self.head.init_weights()
            if isinstance(pretrained, str):
                check_point = torch.load(pretrained)
                head_check_point = {
                    k.replace("output", "fc"): v.squeeze()
                    for k, v in check_point.items() if "output" in k
                }
                self.head.load_state_dict(head_check_point, strict=True)
                logger = get_root_logger()
                logger.info("DLA head load from pretrained model...")
